#include <Functions/IFunction.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/GeoHash.h>

#include <Columns/ColumnArray.h>
#include <Columns/ColumnString.h>
#include <DataTypes/DataTypeArray.h>
#include <DataTypes/DataTypeString.h>

#include <memory>
#include <string>


namespace DB
{

namespace ErrorCodes
{
extern const int LOGICAL_ERROR;
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
extern const int TOO_LARGE_ARRAY_SIZE;
}

namespace
{

class FunctionGeohashesInBox : public IFunction
{
public:
    static constexpr auto name = "geohashesInBox";
    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionGeohashesInBox>(); }

    String getName() const override { return name; }

    size_t getNumberOfArguments() const override { return 5; }

    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
    {
        FunctionArgumentDescriptors args{
            {"longitute_min", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isFloat), nullptr, "Float*"},
            {"latitude_min", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isFloat), nullptr, "Float*"},
            {"longitute_max", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isFloat), nullptr, "Float*"},
            {"latitude_max", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isFloat), nullptr, "Float*"},
            {"precision", static_cast<FunctionArgumentDescriptor::TypeValidator>(&isUInt8), nullptr, "UInt8"}
        };
        validateFunctionArguments(*this, arguments, args);

        if (!(arguments[0].type->equals(*arguments[1].type) &&
              arguments[0].type->equals(*arguments[2].type) &&
              arguments[0].type->equals(*arguments[3].type)))
        {
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
                            "Illegal type of argument of {} all coordinate arguments must have the same type, "
                            "instead they are:{}, {}, {}, {}.", getName(), arguments[0].type->getName(),
                            arguments[1].type->getName(), arguments[2].type->getName(), arguments[3].type->getName());
        }

        return std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>());
    }

    bool useDefaultImplementationForConstants() const override { return true; }

    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }

    template <typename LonAndLatType, typename PrecisionType>
    void execute(
        const IColumn * lon_min_column,
        const IColumn * lat_min_column,
        const IColumn * lon_max_column,
        const IColumn * lat_max_column,
        const IColumn * precision_column,
        ColumnPtr & result,
        size_t input_rows_count) const
    {
        static constexpr size_t max_array_size = 10'000'000;

        const auto * lon_min_const = typeid_cast<const ColumnConst *>(lon_min_column);
        const auto * lat_min_const = typeid_cast<const ColumnConst *>(lat_min_column);
        const auto * lon_max_const = typeid_cast<const ColumnConst *>(lon_max_column);
        const auto * lat_max_const = typeid_cast<const ColumnConst *>(lat_max_column);
        const auto * precision_const = typeid_cast<const ColumnConst *>(precision_column);

        if (lon_min_const)
            lon_min_column = &lon_min_const->getDataColumn();
        if (lat_min_const)
            lat_min_column = &lat_min_const->getDataColumn();
        if (lon_max_const)
            lon_max_column = &lon_max_const->getDataColumn();
        if (lat_max_const)
            lat_max_column = &lat_max_const->getDataColumn();
        if (precision_const)
            precision_column = &precision_const->getDataColumn();

        const auto * lon_min = checkAndGetColumn<ColumnVector<LonAndLatType>>(lon_min_column);
        const auto * lat_min = checkAndGetColumn<ColumnVector<LonAndLatType>>(lat_min_column);
        const auto * lon_max = checkAndGetColumn<ColumnVector<LonAndLatType>>(lon_max_column);
        const auto * lat_max = checkAndGetColumn<ColumnVector<LonAndLatType>>(lat_max_column);
        const auto * precision = checkAndGetColumn<ColumnVector<PrecisionType>>(precision_column);

        if (!lon_min || !lat_min || !lon_max || !lat_max || !precision)
        {
            throw Exception(ErrorCodes::LOGICAL_ERROR, "Unsupported argument types for function {} : {}, {}, {}, {}.",
                            getName(), lon_min_column->getName(),
                            lat_min_column->getName(), lon_max_column->getName(), lat_max_column->getName());
        }

        auto col_res = ColumnArray::create(ColumnString::create());
        ColumnString & res_strings = typeid_cast<ColumnString &>(col_res->getData());
        ColumnArray::Offsets & res_offsets = col_res->getOffsets();
        ColumnString::Chars & res_strings_chars = res_strings.getChars();
        ColumnString::Offsets & res_strings_offsets = res_strings.getOffsets();

        for (size_t row = 0; row < input_rows_count; ++row)
        {
            const Float64 lon_min_value = lon_min->getElement(lon_min_const ? 0 : row);
            const Float64 lat_min_value = lat_min->getElement(lat_min_const ? 0 : row);
            const Float64 lon_max_value = lon_max->getElement(lon_max_const ? 0 : row);
            const Float64 lat_max_value = lat_max->getElement(lat_max_const ? 0 : row);
            const PrecisionType precision_value = precision->getElement(precision_const ? 0 : row);

            const auto prepared_args = geohashesInBoxPrepare(
                lon_min_value, lat_min_value, lon_max_value, lat_max_value,
                precision_value);

            if (prepared_args.items_count > max_array_size)
            {
                throw Exception(ErrorCodes::TOO_LARGE_ARRAY_SIZE, "{} would produce {} array elements, "
                                "which is bigger than the allowed maximum of {}",
                                getName(), prepared_args.items_count, max_array_size);
            }

            res_strings_offsets.reserve(res_strings_offsets.size() + prepared_args.items_count);
            res_strings_chars.resize(res_strings_chars.size() + prepared_args.items_count * prepared_args.precision);
            const auto starting_offset = res_strings_offsets.empty() ? 0 : res_strings_offsets.back();
            char * out = reinterpret_cast<char *>(res_strings_chars.data() + starting_offset);

            // Actually write geohashes into preallocated buffer.
            geohashesInBox(prepared_args, out);

            for (UInt64 i = 1; i <= prepared_args.items_count ; ++i)
                res_strings_offsets.push_back(starting_offset + prepared_args.precision * i);
            res_offsets.push_back(res_offsets.back() + prepared_args.items_count);
        }

        if (!res_strings_offsets.empty() && res_strings_offsets.back() != res_strings_chars.size())
        {
            throw Exception(ErrorCodes::LOGICAL_ERROR, "String column size mismatch (internal logical error)");
        }

        if (!res_offsets.empty() && res_offsets.back() != res_strings.size())
        {
            throw Exception(ErrorCodes::LOGICAL_ERROR, "Array column size mismatch (internal logical error){} != {}",
                            res_offsets.back(), std::to_string(res_strings.size()));
        }

        result = std::move(col_res);
    }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
    {
        const IColumn * lon_min = arguments[0].column.get();
        const IColumn * lat_min = arguments[1].column.get();
        const IColumn * lon_max = arguments[2].column.get();
        const IColumn * lat_max = arguments[3].column.get();
        const IColumn * precision = arguments[4].column.get();
        ColumnPtr res;

        if (checkColumn<ColumnVector<Float32>>(lon_min))
            execute<Float32, UInt8>(lon_min, lat_min, lon_max, lat_max, precision, res, input_rows_count);
        else
            execute<Float64, UInt8>(lon_min, lat_min, lon_max, lat_max, precision, res, input_rows_count);

        return res;
    }
};

}

REGISTER_FUNCTION(GeohashesInBox)
{
    FunctionDocumentation::Description description = R"(
Returns an array of [geohash](https://en.wikipedia.org/wiki/Geohash)-encoded strings of given precision that fall inside and intersect boundaries of given box, essentially a 2D grid flattened into an array.

:::note
All coordinate parameters must be of the same type: either `Float32` or `Float64`.
:::

This function throws an exception if the size of the resulting array exceeds more than 10,000,000 items.
    )";
    FunctionDocumentation::Syntax syntax = "geohashesInBox(longitude_min, latitude_min, longitude_max, latitude_max, precision)";
    FunctionDocumentation::Arguments arguments = {
        {"longitude_min", "Minimum longitude. Range: `[-180°, 180°]`.", {"Float32", "Float64"}},
        {"latitude_min", "Minimum latitude. Range: `[-90°, 90°]`.", {"Float32", "Float64"}},
        {"longitude_max", "Maximum longitude. Range: `[-180°, 180°]`.", {"Float32", "Float64"}},
        {"latitude_max", "Maximum latitude. Range: `[-90°, 90°]`.", {"Float32", "Float64"}},
        {"precision", "Geohash precision. Range: `[1, 12]`.", {"UInt8"}}
    };
    FunctionDocumentation::ReturnedValue returned_value = {
        "Returns an array of precision-long strings of geohash-boxes covering the provided area, or an empty array if the minimum longitude and latitude values aren't less than the corresponding maximum values.",
        {"Array(String)"}
    };
    FunctionDocumentation::Examples examples = {
        {
            "Basic usage",
            "SELECT geohashesInBox(24.48, 40.56, 24.785, 40.81, 4) AS thasos",
            R"(
┌─thasos──────────────────────────────────────┐
│ ['sx1q','sx1r','sx32','sx1w','sx1x','sx38'] │
└─────────────────────────────────────────────┘
            )"
        }
    };
    FunctionDocumentation::IntroducedIn introduced_in = {20, 1};
    FunctionDocumentation::Category category = FunctionDocumentation::Category::Geo;
    FunctionDocumentation documentation = {description, syntax, arguments, returned_value, examples, introduced_in, category};
    factory.registerFunction<FunctionGeohashesInBox>(documentation);
}

}
